# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import datetime
import functools
import logging

from absl import app, flags
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.flow.simple_queue import SimpleQueue
from coin.strategy.hk_hard_arb.balance_printer1 import (ArbInfoPrinter, FungibleArbProducts)
from coin.strategy.hk_hard_arb.pricer import EdgeCalculator
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import get_live_env_cls

FLAGS = flags.FLAGS

flags.DEFINE_string('exchanges', 'Exchanges', '')

flags.DEFINE_string('og_configs', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_integer('exit_after_min', 30, '')

flags.DEFINE_boolean('dry_run', False, '')


class SingleProductSubStrategy(object):
  def __init__(
      self,
      arb_product,
      parent,  # TODO(inkyu): Separate me
      config,
      logger=None):
    self._arb_product = arb_product
    self._product = arb_product.product
    self._native_product_0 = arb_product.native_products[0]
    self._native_product_1 = arb_product.native_products[1]

    self._logger = logger

    # Execution params
    self._edge = config['edge_bp'] * 0.0001
    self._close_edge = config['close_edge_bp'] * 0.0001
    self._lot_size_min = config['lot_size_min']
    self._lot_size_max = config['lot_size_max']

    self._reserve_balance = config['reserve_balance']

    # composite instances
    self._parent = parent
    self._pricer = EdgeCalculator(arb_product, price_det_notional=config['price_det_notional'])

    # state
    self._timestamp_last_order = 0
    self._verbose = True

  def _cancel_all(self):
    return (self._parent.og_0.cancel_product(self._product)
            or self._parent.og_1.cancel_product(self._product))

  def _has_pending_order(self):
    return (self._parent.og_0.has_pending_order(self._product)
            or self._parent.og_1.has_pending_order(self._product))

  def _get_edge(self, position_to_sell):
    if position_to_sell >= self._reserve_balance * 0.6:
      return self._close_edge
    else:
      return self._edge

  def on_position_update(self, timestamp):
    # self._on_book(timestamp, True)
    pass

  def on_book(self, exchange_id, book):
    self._pricer.on_book(exchange_id, book)

    # Wait for 1 sec
    if ((book.timestamp - self._timestamp_last_order) / 1000000) < 1000:
      return

    # Return if there is a pending order
    if self._has_pending_order():
      return

    # Cancel all of the existing orders
    if self._cancel_all():
      return

    if not self._pricer.ready:
      return

    self.agg_agg(book.timestamp)

  def agg_agg(self, timestamp):
    arb_balance = self._parent._balance_printer.get_balance(self._arb_product,
                                                            self._reserve_balance)
    balance_overflow = arb_balance.base_offset

    # Calculate qty to buy / sell
    if balance_overflow > 0:
      sell_qty = self._lot_size_min
      buy_qty = min(self._lot_size_max, max(sell_qty - balance_overflow, 0))
      if buy_qty < self._lot_size_min:
        buy_qty = 0
    elif balance_overflow < 0:
      buy_qty = self._lot_size_min
      sell_qty = min(self._lot_size_max, max(buy_qty + balance_overflow, 0))
      if sell_qty < self._lot_size_min:
        sell_qty = 0
    else:
      return

    # TODO(inkyu): Precision
    # sell_qty_s = float_to_str(sell_qty, self._qty_precision)
    # buy_qty_s = float_to_str(buy_qty, self._qty_precision)

    edge = self._get_edge(arb_balance.base_balances[0])
    if self._pricer.edge_0_1 > edge:
      self._logger.info('%s: 0 -> 1 edge %.8f (%.8f)' %
                        (self._product, self._pricer.edge_0_1, edge))
      self._timestamp_last_order = timestamp
      if arb_balance.base_balances[0] < sell_qty:
        if self._verbose:
          self._logger.warn('exchange0 %s position is too small: (%.8f)' %
                            (self._product.base, arb_balance.base_balances[0]))
          self._verbose = False
        return
      self._verbose = True
      if buy_qty > 0:
        buy_price = self._pricer.ask_exchange1
        self._logger.info('BUY exchange1 %s x %f at %f' % (self._product, buy_qty, buy_price))
        if not FLAGS.dry_run:
          self._parent.og_1.submit(
              self._native_product_1,
              buy_price,  # TODO(bitfinex <- string)
              buy_qty,
              OrderSide.BUY,
              OrderType.LIMIT)
      if sell_qty > 0:
        sell_price = self._pricer.bid_exchange0
        self._logger.info('SELL exchange0 %s x %f at %f' % (self._product, sell_qty, sell_price))
        if not FLAGS.dry_run:
          self._parent.og_0.submit(self._native_product_0,
                                   sell_price,
                                   sell_qty,
                                   OrderSide.SELL,
                                   OrderType.LIMIT)

    edge = self._get_edge(arb_balance.base_balances[1])
    if self._pricer.edge_1_0 > edge:
      self._logger.info('%s: 1 -> 0 edge %.8f (%.8f)' %
                        (self._product, self._pricer.edge_1_0, edge))
      self._timestamp_last_order = timestamp
      if arb_balance.base_balances[1] < sell_qty:
        if self._verbose:
          self._logger.warn('exchange1 %s position is too small: (%.8f)' %
                            (self._product.base, arb_balance.base_balances[1]))
          self._verbose = False
        return
      self._verbose = True
      if buy_qty > 0:
        buy_price = self._pricer.ask_exchange0
        self._logger.info('BUY exchange0 %s x %f at %f' % (self._product, buy_qty, buy_price))
        if not FLAGS.dry_run:
          self._parent.og_0.submit(self._native_product_0,
                                   buy_price,
                                   buy_qty,
                                   OrderSide.BUY,
                                   OrderType.LIMIT)
      if sell_qty > 0:
        sell_price = self._pricer.bid_exchange1
        self._logger.info('SELL exchange1 %s x %f at %f' % (self._product, sell_qty, sell_price))
        if not FLAGS.dry_run:
          self._parent.og_1.submit(self._native_product_1,
                                   sell_price,
                                   sell_qty,
                                   OrderSide.SELL,
                                   OrderType.LIMIT)


class ArbStrategy(object):
  def __init__(self, strategy_params):
    self._logger = logging.getLogger(__name__)

    self._feed = MultiFeedSubsystem()
    self._config = strategy_params['config']

    exchange_0 = strategy_params['exchange_0']
    exchange_1 = strategy_params['exchange_1']
    subreq_0 = SubscriptionRequest('Spot', exchange_0, None)
    subreq_1 = SubscriptionRequest('Spot', exchange_1, None)
    self._og_config_0 = strategy_params['og_config_0']
    self._og_config_1 = strategy_params['og_config_1']

    # TODO(someone): Need some unified metaclass
    self._book_builder_clses = (
        get_flow_book_builder_cls(subreq_0),
        get_flow_book_builder_cls(subreq_1),
    )
    self._live_env_clses = (
        get_live_env_cls(subreq_0),
        get_live_env_cls(subreq_1),
    )
    self._live_env_0 = None
    self._live_env_1 = None

    self._subreq_map = {subreq_0: 0, subreq_1: 1}
    self.og_0 = None
    self.og_1 = None

    self._single_product_sub_strategies = {}

    self._arb_products = None
    self._balance_printer = None

  def subscribe_feeds(self):
    common_config = self._config['common']

    self._arb_products = FungibleArbProducts.two(common_config['quote_currency'],
                                                 self._config['products'].keys(),
                                                 self._book_builder_clses[0].ProductType,
                                                 self._book_builder_clses[1].ProductType)

    for symbol, config in self._config['products'].items():
      config = {**common_config, **config}
      arb_product = self._arb_products.by_symbol(symbol)
      self._single_product_sub_strategies[arb_product.product] = SingleProductSubStrategy(
          arb_product, self, config, logger=self._logger)

    self._feed.subscribe_book(self._arb_products.natives[0].products,
                              functools.partial(self._on_book, 0))
    self._feed.subscribe_book(self._arb_products.natives[1].products,
                              functools.partial(self._on_book, 1))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    self._live_env_0 = self._live_env_clses[0](self._arb_products.natives[0].currencies,
                                               self._arb_products.natives[0].products,
                                               self._og_config_0,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub)
    self._live_env_0.on_order_gateway_reset_callback.append(self.register_order_gateway)
    self._live_env_1 = self._live_env_clses[1](self._arb_products.natives[1].currencies,
                                               self._arb_products.natives[1].products,
                                               self._og_config_1,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub)
    self._live_env_1.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_order_gateway(self, subreq, order_gateway):
    self._logger.info('OG initialized: %s', subreq)
    subreq_id = self._subreq_id_map[subreq]
    if subreq_id == 0:
      self.og_0 = order_gateway
    elif subreq_id == 1:
      self.og_1 = order_gateway

    if self._balance_printer is None and self.og_0 and self.og_1:
      self._balance_printer = ArbInfoPrinter([self.og_0, self.og_1],
                                             self._arb_products.shared.quote_currency,
                                             self._arb_products.shared.base_currencies)
      self._balance_printer.print_balance()

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env_0.start()
    self._live_env_1.start()
    self.print_balance(periodic=True)

  def _on_book(self, subreq_id, book):
    if not self.og_0:
      self._logger.info('OG 0 not ready')
      return
    if not self.og_1:
      self._logger.info('OG 1 not ready')
      return

    product = SharedProduct.FromProduct(book.product)
    self._single_product_sub_strategies[product].on_book(subreq_id, book)

  def print_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_balance()
    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.print_balance, True)


def get_param_set(exchanges, og_configs, config_filename):
  exchange_0, exchange_1 = exchanges.split(',')
  og_config_0, og_config_1 = og_configs.split(',')
  config = load_clean_json(config_filename)
  return {
      'exchange_0': exchange_0,
      'exchange_1': exchange_1,
      'og_config_0': og_config_0,
      'og_config_1': og_config_1,
      'config': config,
  }


def main(argv):
  strategy_params = get_param_set(exchanges=FLAGS.exchanges,
                                  og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('hard_arb')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()
  strategy = ArbStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)  # What the..
  strategy.start(ioloop, queue, flow_sub)

  exit_after = datetime.timedelta(minutes=FLAGS.exit_after_min)
  ioloop.add_timeout(exit_after, ioloop.stop)
  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
